package cc.yihy.spider;

import com.alibaba.fastjson.JSON;
import org.apache.commons.codec.digest.DigestUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.HashOperations;
import org.springframework.data.redis.core.ListOperations;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.SetOperations;
import org.springframework.stereotype.Component;
import us.codecraft.webmagic.Request;
import us.codecraft.webmagic.Task;
import us.codecraft.webmagic.scheduler.DuplicateRemovedScheduler;
import us.codecraft.webmagic.scheduler.MonitorableScheduler;
import us.codecraft.webmagic.scheduler.component.DuplicateRemover;

import javax.annotation.Resource;

@Component
public class RedisDataScheduler extends DuplicateRemovedScheduler implements MonitorableScheduler, DuplicateRemover {

    private static final String QUEUE_PREFIX = "queue_";

    private static final String SET_PREFIX = "set_";

    private static final String ITEM_PREFIX = "item_";
    private static final String HEADER_PREFIX = "header_";

    // inject the actual template
    @Autowired
    private RedisTemplate<String, String> redisTemplate; // inject the template as ListOperations

    @Resource(name = "redisTemplate")
    private SetOperations<String, String> setOps;
    @Resource(name = "redisTemplate")
    private ListOperations<String, String> listOps;

    @Resource(name = "redisTemplate")
    private HashOperations<String, String, String> hashOps;

    public RedisDataScheduler() {
        setDuplicateRemover(this);
    }

    /**
     * 添加集合元素
     *
     * @param request
     * @param task
     * @return
     */
    @Override
    public boolean isDuplicate(Request request, Task task) {
        return setOps.add(getSetKey(task), request.getUrl()).intValue() == 0;

    }


    @Override
    protected void pushWhenNoDuplicate(Request request, Task task) {

        listOps.rightPush(getQueueKey(task), request.getUrl());
        if (request.getExtras() != null||!request.getHeaders().isEmpty()) {
            String field = DigestUtils.shaHex(request.getUrl());
            String value = JSON.toJSONString(request);

            hashOps.put(getItemKey(task), field, value);
        }
    }

    /**
     * 删除一个
     *
     * @param task
     */
    @Override
    public void resetDuplicateCheck(Task task) {


        setOps.remove(getSetKey(task));


    }

    /**
     * 获取一个url,
     *
     * @param task
     * @return
     */
    @Override
    public Request poll(Task task) {

        String url = listOps.leftPop(getQueueKey(task));
        if (url == null) {
            return null;
        }

        String key = getItemKey(task);
        String field = DigestUtils.shaHex(url);
        String value = hashOps.get(key, field);
        if (value != null) {
            Request o = JSON.parseObject(value, Request.class);
            return o;
        }
        Request request = new Request(url);
        return request;

    }


    @Override
    public int getTotalRequestsCount(Task task) {

        Long size = setOps.size(getSetKey(task));
        return size.intValue();
    }

    /**
     * 获取队列中的请求数量
     *
     * @param task
     * @return
     */
    @Override
    public int getLeftRequestsCount(Task task) {

        Long size = listOps.size(getQueueKey(task));
        return size.intValue();
    }


    protected String getSetKey(Task task) {
        return SET_PREFIX + task.getUUID();
    }

    protected String getQueueKey(Task task) {
        return QUEUE_PREFIX + task.getUUID();
    }

    protected String getItemKey(Task task) {
        return ITEM_PREFIX + task.getUUID();
    }
    protected String getHeaderKey(Task task) {
        return HEADER_PREFIX + task.getUUID();
    }
}
